/* 
 * Flv timestamp hack util, (C) 2010 <neko@mikufans.cn> nekofs.
 *
 * compile with:
 *   gcc flvtshack.c -o flvtshack -Wall -Wextra
 */ 

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h> ///for getopt(). i'm lazy =w=
#include "flv.h"

/** flv print functions  */

void flvheader_print(FILE *f, flvhdr *fh){
  fprintf(f, "FLV%u video:%s audio:%s data-offset:%u\n", 
    fh->ver, fh->has_video ? "yes" : "no", fh->has_audio ? "yes" : "no",
    bswap32(fh->doff));
}

void flvtag_print(FILE* f, flvtaghdr *fth, int num){
  fprintf(f, "#%d %s time:%.3f size:%u%c", num,
    fth->type == FLVTAG_AUDIO ? "audio" 
    : fth->type == FLVTAG_VIDEO ? "video"
    : fth->type == FLVTAG_SCRIPT ? "script" : "reserved",
    bswap24(fth->ts)/1e3, bswap24(fth->ds),
    fth->type == FLVTAG_SCRIPT ? '\n' : ' ');
}

void flvaudio_print(FILE *f, flvaudiohdr *fah){
  char *format[16] = {"PCM", "ADPCM", "MP3", "PCM LE",
    NULL, NULL, "Nellymoser", "G.711 A-law", "G.711 mu-law", NULL, "AAC", 
    "Speex", NULL, NULL, NULL, "Device-specific"};
  char *rate[4] = {"5.5k", "11k", "22k", "44k"};
  char *depth[2] = {"8bit", "16bit"};
  char *type[2] = {"mono", "stereo"};

  switch(fah->format){
  case 0:
  case 1:
  case 2:
  case 3:
  case 6:
  case 7:
  case 8:
  case 10:
  case 11:
  case 15:
    fprintf(f, "%s %s %s %s\n", format[fah->format], 
      rate[fah->rate], depth[fah->depth], type[fah->type]); break;
  case 4:
    fprintf(f, "Nellymoser 16k %s mono\n", depth[fah->depth]); break;
  case 5:
    fprintf(f, "Nellymoser 8k %s mono\n", depth[fah->depth]); break;
  case 14:
    fprintf(f, "MP3 8k %s %s\n", depth[fah->depth], type[fah->type]); break;
  default:
    fputc('\n', f);
  }
}

void flvvideo_print(FILE *f, flvvideohdr *fvh){
  char *type[5] = {"keyframe", "interframe", "disposable interframe", 
    "generated keyframe", "command frame"};
  char *codec[7] = {"JPEG", "H263", "FSV1", "VP6", "VP6A", "FSV2", "AVC"};
  if(!(fvh->type >= 1 && fvh->type <= 5)
  || !(fvh->codec >= 1 && fvh->codec <= 7)){
    fputc('\n', f);
    return;
  }
  fprintf(f, "%s %s\n", codec[fvh->codec-1], type[fvh->type-1]);
}

void inject_flvaudio_tag(FILE *f, int ts){
/* NB. This is a legal mp3 frame. But now an empty frame works too, 
  so I comment this out. The following h263pkt is the same case.

  static char mp3frame[104] = "\xff\xfb\x10\xc4";
  static int mp3frame_s = 104; */

  int mp3frame_s = 0;
  u32 prevsize = bswap32(sizeof(flvtaghdr) + sizeof(flvaudiohdr) + mp3frame_s);
  flvtaghdr fth = {
    .type = FLVTAG_AUDIO,
    .ds = bswap24(mp3frame_s + sizeof(flvaudiohdr)),
    .ts = bswap24(ts & 0xffffff),
    .tse = ts >> 24,
    .id = 0,
  };
  flvaudiohdr fah = {
    .format = FLVAUDIO_MP3,
    .rate = FLVAUDIO_44100HZ,
    .depth = FLVAUDIO_16BIT,
    .type = FLVAUDIO_MONO,
  };

  fwrite(&fth, sizeof(fth), 1, f);
  fwrite(&fah, sizeof(fah), 1, f);
  //fwrite(mp3frame, 104, 1, f);
  fwrite(&prevsize, 4, 1, f);
}

void inject_flvvideo_tag(FILE *f, int ts, int keyframe){
/*  static char h263pkt[18] = 
      "\0\0\x84\0\x04\x04\x11\x27\xd7\xd7\xd7\xd7\xff\xde\0\0\0\x1a";
  static int h263pkt_s = 18; */
  int h263pkt_s = 0;
  u32 prevsize = bswap32(sizeof(flvtaghdr) + sizeof(flvvideohdr) + h263pkt_s);
  flvtaghdr fth = {
    .type = FLVTAG_VIDEO,
    .ds = bswap24(h263pkt_s + sizeof(flvvideohdr)),
    .ts = bswap24(ts & 0xffffff),
    .tse = ts >> 24,
    .id = 0,
  };
  flvvideohdr fvh = {
    .type = keyframe ? FLVVIDEO_KEY : FLVVIDEO_INTER,
    .codec = FLVVIDEO_H263,
  };
  fwrite(&fth, sizeof(fth), 1, f);
  fwrite(&fvh, sizeof(fvh), 1, f);
  //fwrite(h263pkt, 18, 1, f);
  fwrite(&prevsize, 4, 1, f);
}

int drop_meta = 1;
char *foutname = NULL;
char *finname = NULL;
int list_mode = 1;
int match_keyframe = 0;
int format_spoofing = 0;
double scaling = 1.;

void usage(){
  fprintf(stderr, 
"Flv timestamp hack util, (C) 2010 <neko@mikufans.cn> nekofs.\n\
  flvtshack [-p] [-s scale] [-o filename] example.flv\n\
Options:\n\
    -h             show this help\n\
    -p             format spoofing by injecting two forged flv tags to \n\
                   file header.\n\
    -s <scale>     time scaling. duration will be increased to scale_int \n\
                   times by linearly scaling every\n\
                   frame's timestamp. useful for bitrate spoofing.\n\
  -o <filename>  output file. `-' for stdout.\n\
\n\
Note\n\
  When neither -p nor -s option is set, a list about information of every flv\n\
  tag will be printed to stdout.\n\
  The onMetaData tag is removed by default.\n\n");
}

int parse_arg(int argc, char **argv){
  int opt;

  while((opt = getopt(argc, argv, "phs:o:")) != -1){
    switch(opt){
    case 'p':
      format_spoofing = 1;
      break;
    case 's':
      scaling = atof(optarg);
      break;
    case 'o':
      foutname = optarg;
      break;
    case 'h':
    default:
      usage();
      exit(0);
    }
  }
  if(optind < argc)
    finname = argv[optind];

  if(finname == NULL
  || strncasecmp(finname + strlen(finname) - 4, ".flv", 4)){
    fprintf(stderr, "** no file name or not a flv file.\n\n");
    usage();
    return -1;
  }

  if(scaling <= 0){
    fprintf(stderr, "** invalid scaling factor.\n\n");
    usage();
    return -1;
  }

  if(format_spoofing || scaling != 1.)
    list_mode = 0;
    
  return 0;
}

int main(int argc, char **argv){
  FILE *fin = NULL;
  FILE *fout = NULL;
  char buf[BUFSIZ];
  flvhdr fh;
  flvtaghdr fth;
  flvaudiohdr fah;
  flvvideohdr fvh;
  size_t n, toread, tagnum = 0;
  int nowrite;

  
  if(parse_arg(argc, argv) < 0)
    return 1;

  if((fin = fopen(finname, "rb")) == NULL){
    fprintf(stderr, "can't open %s for reading\n", finname);
    return 1;
  }

  if(!list_mode){
    if(foutname == NULL){
      foutname = malloc(strlen(finname)+2+1); //XXX I leave free() to system.
      if(foutname == NULL){
        fprintf(stderr, "malloc failed\n");
        fclose(fin);
        return 1;
      }
      strcpy(foutname, finname);
      strcpy(foutname + strlen(finname) - 4, "-1.flv");
    }
    if(!strcmp(foutname, "-"))
      fout = stdout;
    else if((fout = fopen(foutname, "wb")) == NULL){
      fprintf(stderr, "can't open %s for writing\n", foutname);
      return 1;
    }
  }

  fread(&fh, sizeof(fh), 1, fin);

  if(strncmp((char*)fh.sign, "FLV", 3) || !fh.has_video){
    fprintf(stderr, "invalid format\n");
    return 1;
  }

  if(!list_mode)
    fwrite(&fh, sizeof(fh), 1, fout);

  if(list_mode)flvheader_print(stdout, &fh);

  if(list_mode){
    fseek(fin, bswap32(fh.doff) - sizeof(fh), SEEK_CUR);
  }else{
    for(toread = bswap32(fh.doff) - sizeof(fh); toread > 0;){
      n = fread(buf, 1, toread > BUFSIZ ? BUFSIZ : toread, fin);
      fwrite(buf, 1, n, fout);
      toread -= n;
    }
  }

  //PreviousTagSize0
  if(list_mode){
    fseek(fin, 4, SEEK_CUR);
  }else{
    fread(buf, 4, 1, fin);
    fwrite(buf, 4, 1, fout);
    if(format_spoofing){
      inject_flvvideo_tag(fout, 0, 1);
      inject_flvaudio_tag(fout, 0);
    }
  }

  for(;;){
    if(!fread(&fth, sizeof(fth), 1, fin))
      break;

    if(list_mode)flvtag_print(stdout, &fth, tagnum++);
    if(fth.type == FLVTAG_AUDIO){
      fread(&fah, sizeof(fah), 1, fin);
      fseek(fin, -sizeof(fah), SEEK_CUR);
      if(list_mode)flvaudio_print(stdout, &fah);
    }else if(fth.type == FLVTAG_VIDEO){
      fread(&fvh, sizeof(fvh), 1, fin);
      fseek(fin, -sizeof(fvh), SEEK_CUR);
      if(list_mode)flvvideo_print(stdout, &fvh);
    }

    nowrite = 0;
    if(!list_mode && scaling != 1.){
      int ts;
      ts = bswap24(fth.ts) | (fth.tse << 24);

      /* HERE! do something to timestamp */
      //fprintf(stderr, "tag time %.3f -> %.3f\n", fth.ts/1e3, fth.ts/1e3+offset);
      ts = ts * scaling;

      fth.ts = bswap24(ts & 0xffffff);
      fth.tse = ts >> 24;
    }

    if(!list_mode && fth.type == FLVTAG_SCRIPT && drop_meta){
      fread(buf, 13, 1, fin); //XXX
      fseek(fin, -13, SEEK_CUR);
      if(memcmp(buf, "\x02\x00\x0aonMetaData", 13) == 0)
        nowrite = 1;
    }
    
    if(nowrite || list_mode){
      fseek(fin, bswap24(fth.ds)+4, SEEK_CUR);
    }else{
      fwrite(&fth, sizeof(fth), 1, fout);
      for(toread = bswap24(fth.ds); toread > 0;){
        n = fread(buf, 1, toread > BUFSIZ ? BUFSIZ : toread, fin);
        fwrite(buf, n, 1, fout);
        toread -= n;
      }
      fread(buf, 4, 1, fin);
      fwrite(buf, 4, 1, fout);
    }
  }

  fclose(fin);
  if(fout)fclose(fout);
  return 0;
}
